# -*- coding: utf-8 -*-
"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import argparse
import pylynchipsdk as sdk
import numpy as np
import threading
import _thread
import struct
from ctypes import *
from typing import cast
import getopt

import os
import sys
import time
sys.path.append('../../../')
from common.python.plugin_utils import *
import common.python.common as common
import common.python.bufferpool as bufferpool

class decoder_attr:
    def __init__(self) -> None:
        self.image_path = ""
        self.plugin_path = ""
        self.model_path = ""
        self.device_id = 0

def usage():
    print("usage: python3 main.py")
    print("\t -i    image path, eg: ../../../data/yolov3_detect_with_image_sample/test.jpeg")
    print("\t -p    post plugin path, eg: ../../../common/plugin/postprocess/cross_compiled_lib/libYolov3PostPlugin.so")
    print("\t -o    osd plugin path, eg: ../../../common/plugin/osd/lib/libOsdPlugin.so")
    print("\t -l    label file path, eg: ../../../common/plugin/postprocess/yolox/label/label.v3.json")
    print("\t -m    model path, eg: ../../../model/yolov3_b1_h608_w608_c3_iuint8_ofloat16_0000_quant/Net_0")
    print("\t -d    device id, default 0")
    print("\t -h    show help")
    print("\t -?    show help")

def free_to_pool_callback(params):
    data = params[0]
    ret = sdk.lyn_free(data)
    common.error_check(ret, "lyn_free")
    return 0

def set_even(num):
    if num % 2 != 0:
        res = num + 1
    else:
        res = num
    return res

def set_padding_data(video_width, video_heigh, model_width, model_height):
    if (video_width > video_heigh):
        resize_width = model_width
        pad_x = 0
        # resize_height = (video_heigh * model_width / video_width) & ~1
        resize_height = set_even((video_heigh * model_width // video_width))
        pad_y = (model_height - resize_height) // 2
    else:
        resize_height = model_height
        pad_y = 0
        # resize_width = (video_width * model_width / video_heigh) & ~1
        resize_width = set_even((video_width * model_width // video_heigh))
        pad_x = (model_width - resize_width) // 2

    return int(resize_width), int(resize_height), int(pad_x), int(pad_y)

def main():

    # 校验用户设置的 device id
    if not common.check_device_id(attr.device_id):
        return False

    # 创建 context
    context, ret = sdk.lyn_create_context(attr.device_id)
    common.error_check(ret, "lyn_create_context")
    ret = sdk.lyn_register_error_handler(common.default_stream_error_handle, None)
    common.error_check(ret, "lyn_register_error_handler")
    # 创建 stream
    stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")

    ipe_stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")
    apu_stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")
    post_stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")

    # 创建 event
    ipe_event, ret = sdk.lyn_create_event()
    common.error_check(ret, "lyn_create_event")
    apu_event, ret = sdk.lyn_create_event()
    common.error_check(ret, "lyn_create_event")

    label_list, label_info_dev = create_label_list_with_names(attr.label_file)

   # 3.读取文件
    np_jpg = np.fromfile(attr.image_path, dtype=np.byte)
    in_buffer_size = np_jpg.itemsize * np_jpg.size
    np_jpg_ptr = sdk.lyn_numpy_to_ptr(np_jpg)

    # 4.获取图片信息
    para = sdk.lyn_image_dec_para_t()
    para.img_url = attr.image_path
    para.scale = sdk.lyn_scale_t.SCALE_NONE
    para.output_fmt = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
    para.align = True
    info_ori, support, ret = sdk.lyn_image_get_info(para)
    common.error_check(ret, "lyn_image_get_info")

    print("\ninput info_ori:")
    print("width =", info_ori.input.width)
    print("height =", info_ori.input.height)
    print("codec_id =", info_ori.input.codec_id)
    print("format =", info_ori.input.format)

    print("\noutput info_ori:")
    print("width =", info_ori.output.width)
    print("height =", info_ori.output.height)
    print("scale =", info_ori.output.scale)
    print("output_fmt =", info_ori.output.output_fmt)
    print("align =", info_ori.output.align)
    print("predict_buf_size =", info_ori.output.predict_buf_size)

    print("\nsupport =", support)

    # 5.填充解码属性
    attr_jdec = sdk.lyn_image_dec_attr_t()
    attr_jdec.in_buf.size = in_buffer_size
    attr_jdec.info = info_ori
    attr_jdec.out_buf.size = info_ori.output.predict_buf_size

    dec_frame_data, ret = sdk.lyn_malloc(attr_jdec.out_buf.size)
    common.error_check(ret, "dec_frame_data lyn_malloc")
    # np_output = 0
    if (support):
        # 6.1 硬解码
        attr_jdec.in_buf.data, ret = sdk.lyn_malloc(attr_jdec.in_buf.size)
        common.error_check(ret, "lyn_malloc")
        sdk.lyn_memcpy(attr_jdec.in_buf.data, np_jpg_ptr, attr_jdec.in_buf.size, sdk.lyn_memcpy_dir_t.ClientToServer)
        common.error_check(ret, "lyn_memcpy")
        attr_jdec.out_buf.data = dec_frame_data
        ret = sdk.lyn_jpeg_decode_async(stream, attr_jdec)
        common.error_check(ret, "lyn_jpeg_decode_async")
        ret = sdk.lyn_synchronize_stream(stream)
        common.error_check(ret, "lyn_synchronize_stream")
        sdk.lyn_free(attr_jdec.in_buf.data)
    else:
        # 6.2 软解码
        attr_jdec.in_buf.data = np_jpg_ptr
        attr_jdec.out_buf.data = sdk.c_malloc(attr_jdec.out_buf.size)
        ret = sdk.lyn_image_decode_soft(attr_jdec)
        common.error_check(ret, 'lyn_image_decode_soft')
        sdk.lyn_memcpy(dec_frame_data, attr_jdec.out_buf.data, attr_jdec.out_buf.size, sdk.lyn_memcpy_dir_t.ClientToServer)
        sdk.c_free(attr_jdec.out_buf.data)

    # 生成输出文件路径
    file_name, fe_name = os.path.splitext(attr.image_path)
    file_path = file_name + "_py_result.JPEG"
    print("result save to ", file_path)
    if os.path.exists(file_path):
        os.remove(file_path)

    # 设置编码参数
    attr_jenc = sdk.lyn_jenc_attr_t()
    attr_jenc.input_format = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
    attr_jenc.width = info_ori.output.width
    attr_jenc.height = info_ori.output.height
    attr_jenc.jpeg_quality = 80

    info_jenc = sdk.lyn_jenc_info_t()
    info_jenc.params = attr_jenc

    input_frame_size = attr_jenc.width * attr_jenc.height * 3 // 2
    info_jenc.input.size = input_frame_size

    # 加载模型
    model, ret = sdk.lyn_load_model(attr.model_path)
    common.error_check(ret, "lyn_load_model")
    model_desc, ret = sdk.lyn_model_get_desc(model)
    common.error_check(ret, "lyn_model_get_desc")
    batch_size = model_desc.inputTensorAttrArray[0].batchSize
    model_width = model_desc.inputTensorAttrArray[0].dims[2]
    model_height = model_desc.inputTensorAttrArray[0].dims[1]
    ipe_output_size = model_width * model_height * model_desc.inputTensorAttrArray[0].dims[3]
    apu_output_size = model_desc.outputDataLen

    # 创建 IPE 相关结构体
    ipe_input_desc, ret = sdk.lyn_ipe_create_pic_desc()
    common.error_check(ret, "lyn_ipe_create_pic_desc ipe_input_desc")
    ipe_output_desc, ret = sdk.lyn_ipe_create_pic_desc()
    common.error_check(ret, "lyn_ipe_create_pic_desc ipe_output_desc")
    ipe_config, ret = sdk.lyn_ipe_create_config_desc()
    common.error_check(ret, "lyn_ipe_create_config_desc")
    ret = sdk.lyn_ipe_reset_pic_desc(ipe_input_desc)
    common.error_check(ret, "lyn_ipe_reset_pic_desc")
    ret = sdk.lyn_ipe_reset_pic_desc(ipe_output_desc)
    common.error_check(ret, "lyn_ipe_reset_pic_desc")
    ret = sdk.lyn_ipe_reset_config_desc(ipe_config)
    common.error_check(ret, "lyn_ipe_reset_config_desc")
    resize_width, resize_height, pad_x, pad_y = set_padding_data(info_ori.output.width, info_ori.output.height, model_width, model_height)


    # 加载 Plugin
    osd_plugin, ret = sdk.lyn_plugin_register(attr.osd_plugin)
    common.error_check(ret, "lyn_plugin_register")
    post_plugin, ret = sdk.lyn_plugin_register(attr.plugin_path)
    common.error_check(ret, "lyn_plugin_register")
    boxes_info, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
    common.error_check(ret, "malloc")


    start = time.time()
    # 设置IPE参数，并进行IPE处理
    ipe_in_data = dec_frame_data
    ipe_out_data, ret = sdk.lyn_malloc(ipe_output_size)
    common.error_check(ret, "ipe_out_data lyn_malloc")
    ret = sdk.lyn_ipe_set_input_pic_desc(ipe_input_desc, ipe_in_data, info_ori.output.width,
                                            info_ori.output.height, info_ori.output.output_fmt) # set ipe input
    common.error_check(ret, "lyn_ipe_set_input_pic_desc")
    ret = sdk.lyn_ipe_set_output_pic_data(ipe_output_desc, ipe_out_data) # set ipe output
    common.error_check(ret, "lyn_ipe_set_output_pic_data")
    ret = sdk.lyn_ipe_set_resize_config(ipe_config, resize_width, resize_height) # resize to model require
    common.error_check(ret, "lyn_ipe_set_resize_config")
    ret = sdk.lyn_ipe_set_pad_config(ipe_config, pad_y, pad_x, pad_y, pad_x, 114, 114, 114)
    common.error_check(ret, "lyn_ipe_set_pad_config")
    ret = sdk.lyn_ipe_set_c2c_config(ipe_config, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0)
    common.error_check(ret, "lyn_ipe_set_c2c_config")
    ret = sdk.lyn_ipe_cal_output_pic_desc(ipe_output_desc, ipe_input_desc, ipe_config, 0)
    common.error_check(ret, "lyn_ipe_cal_output_pic_desc")
    ret = sdk.lyn_ipe_process_async(ipe_stream, ipe_input_desc, ipe_output_desc, ipe_config)
    common.error_check(ret, "lyn_ipe_process_async")

    # 等待IPE处理完成
    ret = sdk.lyn_record_event(ipe_stream, ipe_event)
    common.error_check(ret, "lyn_record_event")
    ret = sdk.lyn_stream_wait_event(apu_stream, ipe_event)
    common.error_check(ret, "lyn_stream_wait_event")

    # apu 推理
    apu_output, ret = sdk.lyn_malloc(apu_output_size)
    common.error_check(ret, "apu_output lyn_malloc")
    ret = sdk.lyn_execute_model_async(apu_stream, model, ipe_out_data, apu_output, batch_size)
    common.error_check(ret, "lyn_execute_model_async")
    ret = sdk.lyn_stream_add_callback(apu_stream, free_to_pool_callback, [ipe_out_data])
    common.error_check(ret, "lyn_stream_add_callback")

    # 在 APU 推理完之后，调用 Plugin 进行后处理
    ret = sdk.lyn_record_event(apu_stream, apu_event)
    common.error_check(ret, "lyn_record_event")
    ret = sdk.lyn_stream_wait_event(post_stream, apu_event)
    common.error_check(ret, "lyn_stream_wait_event")
    format = int(sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)
    pythonapi.PyCapsule_GetPointer.restype = c_void_p
    pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]
    apu_data_ptr = pythonapi.PyCapsule_GetPointer(apu_output, None)
    boxes_info_ptr = pythonapi.PyCapsule_GetPointer(boxes_info, None)
    frame_data_ptr = pythonapi.PyCapsule_GetPointer(dec_frame_data, None)
    # post_para = struct.pack('2P', apu_data_ptr, boxes_info)
    post_para = YoloxPostProcessInfo_t()
    post_para.height = model_height
    post_para.width = model_width
    post_para.ori_height = info_ori.output.height
    post_para.ori_width = info_ori.output.width
    post_para.score_threshold = 0.25
    post_para.nms_threshold = 0.45
    post_para.nms_top_k = 500
    post_para.is_pad_resize = 1
    post_para.output_tensor = apu_data_ptr
    post_para.class_num = label_list.labelNum
    post_para.labelList = label_list
    post_para.boxesInfo = boxes_info_ptr
    post_para.anchorSize = 1

    yolox_data = ctypes.string_at(
        ctypes.addressof(post_para), ctypes.sizeof(post_para)
    )

    ret = sdk.lyn_plugin_run_async(
        post_stream,
        post_plugin,
        "lynYoloxPostProcess",
        yolox_data,
        ctypes.sizeof(YoloxPostProcessInfo_t),
    )

    common.error_check(ret, "lyn_plugin_run_async")

    osd_para = struct.pack(
        "1P3i1P4i",
        boxes_info_ptr,
        info_ori.output.width,
        info_ori.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
        frame_data_ptr,
        24,
        4,
        2,
        4,
    )

    common.error_check(
        sdk.lyn_plugin_run_async(
            post_stream,
            osd_plugin,
            "lynDrawBoxAndText",
            osd_para,
            len(osd_para),
        )
    )

    ret = sdk.lyn_stream_add_callback(post_stream, free_to_pool_callback, [apu_output])
    common.error_check(ret, "lyn_stream_add_callback")

    info_jenc.input.data = dec_frame_data
    info_jenc.output.size = input_frame_size
    info_jenc.output.data, ret = sdk.lyn_malloc(info_jenc.output.size)
    common.error_check(ret, "lyn_malloc")

    # async encode
    ret = sdk.lyn_jpeg_encode_async(post_stream, info_jenc)
    common.error_check(ret, "lyn_jpeg_encode_async")

    # wait finish
    ret = sdk.lyn_synchronize_stream(post_stream)
    common.error_check(ret, "lyn_synchronize_stream")

    #dump encode image
    packet = sdk.lyn_packet_t()
    packet.data = info_jenc.output.data
    packet.size = info_jenc.output.size
    outdata_size, ret = sdk.lyn_enc_get_remote_packet_valid_size(packet)
    common.error_check(ret, 'lyn_enc_get_remote_packet_valid_size')
    outdata = np.zeros(outdata_size, dtype=np.uint8)
    outdata_ptr = sdk.lyn_numpy_to_ptr(outdata)
    ret = sdk.lyn_memcpy(outdata_ptr, packet.data, outdata_size, sdk.lyn_memcpy_dir_t.ServerToClient)
    common.error_check(ret, "lyn_memcpy")

    with open(file_path, "wb") as f:
        f.write(outdata.tobytes())
    end = time.time()

    # 等待流处理完
    ret = sdk.lyn_synchronize_stream(ipe_stream)
    common.error_check(ret, "lyn_synchronize_stream")
    ret = sdk.lyn_synchronize_stream(apu_stream)
    common.error_check(ret, "lyn_synchronize_stream")
    ret = sdk.lyn_destroy_stream(ipe_stream)
    common.error_check(ret, "lyn_destroy_stream")
    ret = sdk.lyn_destroy_stream(apu_stream)
    common.error_check(ret, "lyn_destroy_stream")
    ret = sdk.lyn_destroy_stream(stream)
    common.error_check(ret, "lyn_destroy_stream")
    ret = sdk.lyn_destroy_stream(post_stream)
    common.error_check(ret, "lyn_destroy_stream")
    # destory ipe desc and config
    ret = sdk.lyn_ipe_destroy_pic_desc(ipe_input_desc)
    common.error_check(ret, "lyn_ipe_destroy_pic_desc")
    ret = sdk.lyn_ipe_destroy_pic_desc(ipe_output_desc)
    common.error_check(ret, "lyn_ipe_destroy_pic_desc")
    ret = sdk.lyn_ipe_destroy_config_desc(ipe_config)
    common.error_check(ret, "lyn_ipe_destroy_config_desc")

    sdk.lyn_free(info_jenc.output.data)
    sdk.lyn_free(dec_frame_data)
    sdk.lyn_free(boxes_info)
    sdk.lyn_free(label_info_dev)

    # 卸载模型
    ret = sdk.lyn_plugin_unregister(post_plugin)
    common.error_check(ret, "lyn_plugin_unregister")
    ret = sdk.lyn_plugin_unregister(osd_plugin)
    common.error_check(ret, "lyn_plugin_unregister")
    ret = sdk.lyn_unload_model(model)
    common.error_check(ret, "lyn_unload_model")
    # 销毁 context
    ret = sdk.lyn_destroy_context(context)
    common.error_check(ret, "lyn_destroy_context")


if __name__ == "__main__":
    start = time.time()

    argv = sys.argv[1:]
    try:
        opts, args = getopt.getopt(argv, "i:p:o:l:m:d:h?")
    except getopt.GetoptError:
        usage()
        sys.exit(0)

    attr = decoder_attr()

    for opt, arg in opts:
        if opt in("-h", "-?"):
            usage()
            sys.exit(0)
        elif opt == "-i":
            attr.image_path = arg
        elif opt == "-p":
            attr.plugin_path = arg
        elif opt == "-o":
            attr.osd_plugin = arg
        elif opt == "-l":
            attr.label_file = arg
        elif opt == "-m":
            attr.model_path = arg
        elif opt == "-d":
            attr.device_id = int(arg)
        else:
            usage()
            sys.exit(0)

    if (attr.image_path == ""):
        print("Plase input video path")
        usage()
        sys.exit(0)
    elif (attr.plugin_path == ""):
        print("Plase input plugin path")
        usage()
        sys.exit(0)
    elif (attr.model_path == ""):
        print("Plase input model path")
        usage()
        sys.exit(0)

    main()
    end = time.time()
    print("need time: ", end - start)